import hqyj_mqtt
import queue
import base64
import numpy as np
import cv2
"""
通过 MQTT 接收 Base64 编码的图像数据，将其转换为 OpenCV 可处理的图像格式，并支持实时显示和保存


这是一个用于捕获和保存通过MQTT接收的图像的实用脚本。

功能：
使用与main.py相同的b64_to_np函数解码base64图像。
使用OpenCV在窗口中显示图像。
按下's'键时将图像保存为PNG文件，文件名递增（如1.png、2.png）。
按下'q'键退出。
连接到相同的MQTT服务器（127.0.0.1:21883），订阅相同主题（bb）。

用途：
可能用于调试或数据收集，允许用户从3D场景中捕获快照以进行分析或测试。
"""


"""将 MQTT 消息中携带的 Base64 编码图像数据，转换为 OpenCV 能直接处理的 BGR 格式图像矩阵（numpy.ndarray）"""
def b64_to_np(image):
    # 此时的image还是一个字典，我们需要的是字典的值，字典的值是base64格式的图像数据
    img_data = base64.b64decode(image['image'])   #提取Base64数据 并使用Base64解码

    # 将字节流转换numpy数组
    img_np = np.frombuffer(img_data, dtype=np.uint8)

    # 使用opencv读取该数组
    img = cv2.imdecode(img_np, cv2.IMREAD_COLOR)  # 使用cv2.imdecode将NumPy数组解码为OpenCV的图像矩阵（BGR格式）

    return img


if __name__ == '__main__':

    q_mqtt_data = queue.Queue(5)
    i = 1

    # 1. 构建mqtt客户端，连接mqtt服务器，方便和3D场景通信
    mqtt_client = hqyj_mqtt.MQTTClient('127.0.0.1', 21883, 'bb', 'aa', 60)

    while True:
        image = q_mqtt_data.get()

        if 'image' in image:
            # 将接收到的消息解析为opencv处理能用的格式
            img = b64_to_np(image)
            cv2.imshow('img', img)
            key = cv2.waitKey(1)
            # print(key)
            if key == ord('q'):
                break
            elif key == ord('s'):
                cv2.imwrite(f'{i}.png', img)
                i += 1
                print('save successful')
            else:
                continue